import os
from model_trainor import train_and_evaluate_model
from sample_predictor import predict_sample

# 配置文件路径
MODEL_PATH = "models/codebert-base"
ASM_DIR = "samples/subtrain/selectedpro/train"  # ASM训练数据目录
LABEL_FILE = "samples/selectedLabels.csv"
CACHE_FILE = "customs/codebert_selectedpro_cache.joblib"  # 特征缓存文件
CUSTOM_FILE = "customs/codebert_selectedpro_classifier.joblib"  # 训练完成的模型文件
# 可选：SVM/LinearSVM/LogisticRegression/XGBoost/MLP
SELECTED_ALGORITHM = ""  # 选择的分类算法（空字符串表示自动选择）

if __name__ == "__main__":
    # 训练并评估模型
    custom = train_and_evaluate_model(MODEL_PATH, ASM_DIR, LABEL_FILE, CACHE_FILE, CUSTOM_FILE, SELECTED_ALGORITHM)

    # 示例：预测样本
    sample_path = "samples/subtrain/selectedpro/test/1eJx34l8pcAFvMuOwrjB.asm"
    if os.path.exists(sample_path):
        print(f"\n使用样本进行预测: {sample_path}")
        prediction, probability = predict_sample(MODEL_PATH, sample_path, custom)

        print(f"  预测结果: {prediction}")
        print("  类别概率分布:")
        for cls, prob in enumerate(probability, start=1):
            print(f"    类别 {cls}: {prob*100:.2f}%")